import os
import numpy as np
from scipy.spatial import cKDTree


class getworklist(object):

    def __init__(self, **kwargs):
        self.wd = kwargs.get('wd', os.getcwd())

    def neighborElements(self, **kwargs):
        elementArray = kwargs.get('element', None)
        filename = kwargs.get('filename', None)
        fmt = kwargs.get('fmt', 'csv')

        numElements = elementArray.shape[0]

        with open(os.path.join(self.wd, '{}.{}'.format(filename, fmt)), 'ab') as f:
            f.seek(0)
            f.truncate(0)
            for i in range(numElements):
                tempWorkList = np.array([])
                nodeList = elementArray[i]
                testBoolean = np.array([elementArray == n for n in nodeList])
                for j in range(testBoolean.shape[2]):
                    tempList = np.where(testBoolean[j])[0]
                    if len(tempList) != 0:
                        tempWorkList = np.concatenate((tempWorkList, tempList), axis=0)
                    j += 1

                tempWorkList = np.unique(tempWorkList)
                tempWorkList = tempWorkList + 1
                tempWorkList = np.delete(tempWorkList, np.where(tempWorkList == i + 1))
                tempWorkList = tempWorkList.reshape((1, -1))
                np.savetxt(f, tempWorkList, fmt='%d', delimiter='\t')
                i += 1

    def neighborElementsKDTREE(self, **kwargs):
        k = kwargs.get('k', 8)
        centroidArray = kwargs.get('ec', None)
        filename = kwargs.get('filename', None)
        fmt = kwargs.get('fmt', 'csv')

        tree = cKDTree(centroidArray)
        index = tree.query(centroidArray, k + 1)[1]
        index = index[:, 1:] + 1
        index = np.sort(index)
        np.savetxt(os.path.join(self.wd, '{}.{}'.format(filename, fmt)), index, fmt='%d',
                   delimiter='\t')

    def neighborNodesKDTREE(self, **kwargs):
        k = kwargs.get('k', 8)
        nodeArray = kwargs.get('node', None)
        filename = kwargs.get('filename', None)
        fmt = kwargs.get('fmt', 'csv')

        tree = cKDTree(nodeArray)
        index = tree.query(nodeArray, k + 1)[1]
        index = index[:, 1:] + 1
        index = np.sort(index)
        np.savetxt(os.path.join(self.wd, '{}.{}'.format(filename, fmt)), index, fmt='%d',
                   delimiter='\t')

    def sameElesetE(self, **kwargs):

        elesetArray = kwargs.get('eleset', None)
        filename = kwargs.get('filename', None)
        fmt = kwargs.get('fmt', 'csv')

        numElements = elesetArray.size

        with open(os.path.join(self.wd, '{}.{}'.format(filename, fmt)), 'ab') as f:
            f.seek(0)
            f.truncate(0)
            for i in range(numElements):
                tempElesetNum = elesetArray[i]
                tb = elesetArray == tempElesetNum
                tempWorkList = np.where(tb)[0]
                tempWorkList = tempWorkList + 1
                tempWorkList = np.delete(tempWorkList, np.where(tempWorkList == i + 1))
                tempWorkList = tempWorkList.reshape((1, -1))
                np.savetxt(f, tempWorkList, fmt='%d', delimiter='\t')
                i += 1

    def sameElesetN(self, **kwargs):
        elesetArray = kwargs.get('eleset', None)
        elementArray = kwargs.get('element', None)
        filename = kwargs.get('filename', None)
        fmt = kwargs.get('fmt', 'csv')

        numNodes = np.unique(elementArray.flatten()).size
        elesetNum = np.unique(elesetArray.flatten())
        numEleset = elesetNum.size

        f = open(os.path.join(self.wd, '{}_TEMP_EN.{}'.format(filename, fmt)), 'ab')
        f.seek(0)
        f.truncate(0)
        for i in elesetNum:
            tempElementNum = np.where(elesetArray == i)[0]
            tempNode = np.array(())
            for j in tempElementNum:
                tempNode = np.append(tempNode, elementArray[j])
            tempNode = np.unique(tempNode).astype(int)
            tempNode = tempNode.reshape((1, -1))
            np.savetxt(f, tempNode, fmt='%d', delimiter='\t')
        f.close()

        bn = np.array(())
        f = open(os.path.join(self.wd, '{}_TEMP_EN.{}'.format(filename, fmt)), 'r')
        tempEN = f.readlines()
        f.close()
        for i in range(numEleset - 1):
            for j in range(i + 1, numEleset):
                temp1 = np.fromstring(tempEN[i], dtype=int, sep='\t')
                temp2 = np.fromstring(tempEN[j], dtype=int, sep='\t')
                tb = np.in1d(temp1, temp2)
                if tb.any():
                    bn = np.append(bn, temp1[np.where(tb)[0]])
                else:
                    pass

        bn = np.sort(np.unique(bn).astype(int))
        np.savetxt(os.path.join(self.wd, '{}_BN.{}'.format(filename, fmt)), bn, delimiter='\t', fmt='%d')

        f = open(os.path.join(self.wd, '{}_TEMP_ENNB.{}'.format(filename, fmt)), 'ab')
        f.seek(0)
        f.truncate(0)

        ne = np.zeros(numNodes).astype(int)

        for i in range(numEleset):
            temp1 = np.fromstring(tempEN[i], dtype=int, sep='\t')
            tb = np.in1d(temp1, bn)
            if tb.all():
                temp1 = np.array([0])
            elif tb.any():
                temp1 = temp1[np.where(~tb)[0]]

            temp1 = temp1.flatten()
            if temp1.any() == 0:
                pass
            else:
                for j in temp1:
                    ne[j - 1] = elesetNum[i]

            temp1 = temp1.reshape((1, -1))
            np.savetxt(f, temp1, delimiter='\t', fmt='%d')

        f.close()
        np.savetxt(os.path.join(self.wd, '{}_TEMP_NE.{}'.format(filename, fmt)), ne, delimiter='\t', fmt='%d')

        f = open(os.path.join(self.wd, '{}.{}'.format(filename, fmt)), 'ab')
        f.seek(0)
        f.truncate(0)
        for i in range(numNodes):
            tempElesetNum = ne[i]
            if tempElesetNum == 0:
                np.savetxt(f, np.array([0]), delimiter='\t', fmt='%d')
            else:
                tempWorkList = np.where(ne == tempElesetNum)[0] + 1
                tempWorkList = np.delete(tempWorkList, np.where(tempWorkList == i + 1))
                tempWorkList = tempWorkList.reshape((1, -1))
                np.savetxt(f, tempWorkList, fmt='%d', delimiter='\t')
            i += 1
        f.close()

    def findBoundary(self, **kwargs):
        """
        :param kwargs:
        es = ELEMENT SET ARRAY
        ne = CAN BE '_io.TextIOWrapper' (f = open(...)) OR 'list' (f.readlines())
        :return:
        """
        es = kwargs.get('eleset', None)
        se = kwargs.get('sameEleset', None)
        filename = kwargs.get('filename', None)
        fmt = kwargs.get('fmt', 'csv')
        sep = kwargs.get('sep', '\t')

        if not isinstance(se, list):
            se = se.readlines()

        numElements = es.size
        elementNum = np.arange(numElements) + 1
        elesetNum = np.unique(es)

        be = np.array(())
        for i in elesetNum:
            tempElementNum = elementNum[np.where(es == i)]
            for j in tempElementNum:
                tempNE = np.fromstring(se[j - 1], dtype=int, sep=sep)
                tb = np.in1d(tempNE, tempElementNum)
                if not tb.all():
                    be = np.append(be, j)
                else:
                    pass

        be = np.sort(be.astype(int))
        np.savetxt(os.path.join(self.wd, '{}.{}'.format(filename, fmt)), be, delimiter='\t', fmt='%d')
        return be

    def elementCentroid(self, **kwargs):
        elementArray = kwargs.get('element', None)
        nodeArray = kwargs.get('node', None)
        filename = kwargs.get('filename', None)
        fmt = kwargs.get('fmt', 'csv')

        nodeDimension = nodeArray.shape[1]
        numElements = elementArray.shape[0]

        centroidCoord = np.empty((0, nodeDimension))
        for i in range(numElements):
            nodeList = elementArray[i]
            tempCentroidCoord = np.zeros((1, nodeDimension))
            numNodesElement = nodeList.shape[0]
            for n in nodeList:
                tempNodeCoord = nodeArray[n - 1, :]
                tempCentroidCoord += tempNodeCoord
            tempCentroidCoord = tempCentroidCoord / numNodesElement
            centroidCoord = np.concatenate((centroidCoord, tempCentroidCoord))

        np.savetxt(os.path.join(self.wd, '{}.{}'.format(filename, fmt)), centroidCoord, fmt='%1.5f',
                   delimiter='\t')

    @staticmethod
    def loadData(**kwargs):

        _wd = kwargs.get('wd', os.getcwd())
        _filename = kwargs.get('filename', None)
        _fmt = kwargs.get('fmt', 'csv')
        _dtype = kwargs.get('dtype', float)
        _usecols = kwargs.get('usecols', None)
        _delimiter = kwargs.get('delimiter', '\t')
        _comments = kwargs.get('comments', '#')
        _skiprows = kwargs.get('skiprows', 0)

        _data = np.loadtxt(os.path.join(_wd, '{}.{}'.format(_filename, _fmt)), delimiter=_delimiter, usecols=_usecols,
                           skiprows=_skiprows, dtype=_dtype, comments=_comments)

        return _data
